/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 */

#include <iostream>
#include <vector>
#include <fstream>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <fcntl.h>

#include "acl/acl.h"
#include "aclnn_sqrt.h"

// 宏定义日志与错误检查
#define SUCCESS 0
#define FAILED 1

#define INFO_LOG(fmt, args...) fprintf(stdout, "[INFO]  " fmt "\n", ##args)
#define ERROR_LOG(fmt, args...) fprintf(stderr, "[ERROR]  " fmt "\n", ##args)

#define CHECK_RET(cond, return_expr) \
    do { if (!(cond)) { return_expr; } } while (0)


// ======================= 工具函数区域 =========================

/**
 * @brief 读取二进制文件内容到 buffer
 */
bool ReadFile(const std::string &filePath, size_t &fileSize, void *buffer, size_t bufferSize)
{
    struct stat sBuf;
    if (stat(filePath.c_str(), &sBuf) == -1 || !S_ISREG(sBuf.st_mode)) {
        ERROR_LOG("Invalid file: %s", filePath.c_str());
        return false;
    }

    std::ifstream file(filePath, std::ios::binary);
    if (!file.is_open()) {
        ERROR_LOG("Open file failed: %s", filePath.c_str());
        return false;
    }

    file.read(static_cast<char *>(buffer), bufferSize);
    fileSize = file.gcount();  // 实际读取的大小
    file.close();

    if (fileSize == 0 || fileSize > bufferSize) {
        ERROR_LOG("Invalid read size from file.");
        return false;
    }
    return true;
}

/**
 * @brief 写数据到二进制文件
 */
bool WriteFile(const std::string &filePath, const void *buffer, size_t size)
{
    if (!buffer) {
        ERROR_LOG("Invalid buffer");
        return false;
    }

    int fd = open(filePath.c_str(), O_RDWR | O_CREAT | O_TRUNC, S_IRUSR | S_IWUSR);
    if (fd < 0) {
        ERROR_LOG("Open file failed: %s", filePath.c_str());
        return false;
    }

    auto written = write(fd, buffer, size);
    close(fd);

    if (written != size) {
        ERROR_LOG("Write failed: only wrote %ld of %ld bytes", written, size);
        return false;
    }

    return true;
}

/**
 * @brief 计算形状数组对应的总元素数
 */
int64_t GetShapeSize(const std::vector<int64_t> &shape)
{
    int64_t size = 1;
    for (auto dim : shape) {
        size *= dim;
    }
    return size;
}

/**
 * @brief 初始化 ACL device 和 stream
 */
int Init(int32_t deviceId, aclrtStream *stream)
{
    CHECK_RET(aclInit(nullptr) == ACL_SUCCESS, ERROR_LOG("aclInit failed"); return FAILED);
    CHECK_RET(aclrtSetDevice(deviceId) == ACL_SUCCESS, ERROR_LOG("Set device failed"); return FAILED);
    CHECK_RET(aclrtCreateStream(stream) == ACL_SUCCESS, ERROR_LOG("Create stream failed"); return FAILED);
    return SUCCESS;
}

/**
 * @brief 构建 ACL Tensor，包括 device 分配和拷贝
 */
template <typename T>
int CreateAclTensor(const std::vector<T> &hostData, const std::vector<int64_t> &shape, void **deviceAddr,
                    aclDataType dataType, aclTensor **tensor)
{
    size_t size = GetShapeSize(shape) * sizeof(T);
    CHECK_RET(aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST) == ACL_SUCCESS,
              ERROR_LOG("Malloc failed"); return FAILED);
    CHECK_RET(aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE) == ACL_SUCCESS,
              ERROR_LOG("Memcpy failed"); return FAILED);

    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType,
                              nullptr, 0, ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr);
    return SUCCESS;
}

// ======================= 主函数入口 =========================

int main()
{
    int32_t deviceId = 0;
    aclrtStream stream;
    if (Init(deviceId, &stream) != SUCCESS) return FAILED;

    std::vector<int64_t> shape = {1024, 1024};
    size_t dataCount = GetShapeSize(shape);
    size_t dataSize = dataCount * sizeof(aclFloat16);

    std::vector<aclFloat16> inputData(dataCount);
    std::vector<aclFloat16> outputData(dataCount);

    // 读取输入
    size_t fileSize = 0;
    if (!ReadFile("../input/input_x.bin", fileSize, inputData.data(), dataSize)) return FAILED;
    INFO_LOG("Input read successfully.");

    // 创建输入输出 tensor
    void *inputDev = nullptr, *outputDev = nullptr;
    aclTensor *inputTensor = nullptr, *outputTensor = nullptr;

    if (CreateAclTensor(inputData, shape, &inputDev, ACL_FLOAT16, &inputTensor) != SUCCESS) return FAILED;
    if (CreateAclTensor(outputData, shape, &outputDev, ACL_FLOAT16, &outputTensor) != SUCCESS) return FAILED;

    // 获取 workspace
    void *workspace = nullptr;
    uint64_t workspaceSize = 0;
    aclOpExecutor *executor = nullptr;

    if (aclnnSqrtGetWorkspaceSize(inputTensor, outputTensor, &workspaceSize, &executor) != ACL_SUCCESS) {
        ERROR_LOG("Get workspace size failed");
        return FAILED;
    }

    if (workspaceSize > 0 && aclrtMalloc(&workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST) != ACL_SUCCESS) {
        ERROR_LOG("Malloc workspace failed");
        return FAILED;
    }

    // 执行 sqrt 算子
    if (aclnnSqrt(workspace, workspaceSize, executor, stream) != ACL_SUCCESS) {
        ERROR_LOG("aclnnSqrt execution failed");
        return FAILED;
    }

    // 同步等待执行结束
    if (aclrtSynchronizeStream(stream) != ACL_SUCCESS) {
        ERROR_LOG("Stream sync failed");
        return FAILED;
    }

    // 复制输出数据到 host
    if (aclrtMemcpy(outputData.data(), dataSize, outputDev, dataSize, ACL_MEMCPY_DEVICE_TO_HOST) != ACL_SUCCESS) {
        ERROR_LOG("Memcpy result failed");
        return FAILED;
    }

    // 写出推理结果
    if (!WriteFile("../output/output_y.bin", outputData.data(), dataSize)) return FAILED;
    INFO_LOG("Output written successfully.");

    // 资源释放
    aclDestroyTensor(inputTensor);
    aclDestroyTensor(outputTensor);
    aclrtFree(inputDev);
    aclrtFree(outputDev);
    if (workspace) aclrtFree(workspace);
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();

    return SUCCESS;
}
